/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at

* http://www.apache.org/licenses/LICENSE-2.0

* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.

* File sample_process.cpp
* Description: handle acl resource
*/
#include "retina/object_detect.h"
#include <cstdint>
#include <fstream>


using namespace std;

namespace {

const static std::vector<std::string> labels_list = {"tank", "board", "box"};
//Inferential output dataset subscript 0 unit is detection box information data

const uint32_t kScoreBufId = 1;
const uint32_t kClassBufId = 2;
const uint32_t kBBoxBufId = 0;

const size_t kModelOutputBoxNum = 15162;
const float kScoreThreshold = 0.1;
const float kNMSThreshold = 0.2;
const uint32_t kLineSolid = 1;

const double kFountScale = 0.4;
const cv::Scalar kFontColor(0, 0, 255);
const uint32_t kLabelOffset = 11;
//Each field subscript in the box message
enum BBoxIndex { TOPLEFTX = 0, TOPLEFTY, BOTTOMRIGHTX, BOTTOMRIGHTY, SCORE, LABEL };
const vector <cv::Scalar> kColors{
            cv::Scalar(237, 149, 100), cv::Scalar(0, 215, 255), cv::Scalar(50, 205, 50),
            cv::Scalar(139, 85, 26)};
}

RetinaObjectDetect::RetinaObjectDetect(const char* modelPath, uint32_t modelWidth,
                           uint32_t modelHeight)
:deviceId_(0), modelWidth_(modelWidth), modelHeight_(modelHeight), imageDataBuf_(nullptr), imageInfoBuf_(nullptr), isInited_(false){
    modelPath_ = modelPath;
    imageDataSize_ = RGBU8_IMAGE_SIZE(modelWidth, modelHeight);
}

RetinaObjectDetect::~RetinaObjectDetect() {
    DestroyResource();
}

Result RetinaObjectDetect::InitResource() {
    // ACL init
    const char *aclConfigPath = "../src/acl.json";
    aclError ret = aclInit(aclConfigPath);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("Acl init failed");
        return FAILED;
    }
    INFO_LOG("Acl init success");

    // open device
    ret = aclrtSetDevice(deviceId_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("Acl open device %d failed", deviceId_);
        return FAILED;
    }
    INFO_LOG("Open device %d success", deviceId_);

    aclrtContext ctx;
    ret = aclrtCreateContext(&ctx, 0);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl create context failed");
        return FAILED;
    }
    aclrtSetCurrentContext(ctx);
    INFO_LOG("create context success");

    ret = aclrtCreateStream(&stream_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl create stream failed");
        return FAILED;
    }
    INFO_LOG("create stream success");

    //Gets whether the current application is running on host or Device
    ret = aclrtGetRunMode(&runMode_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl get run mode failed");
        return FAILED;
    }
    INFO_LOG("GetRunMode success");
    return SUCCESS;
}

Result RetinaObjectDetect::InitModel(const char* omModelPath) {
    Result ret = model_.LoadModelFromFileWithMem(omModelPath);
    if (ret != SUCCESS) {
        ERROR_LOG("execute LoadModelFromFileWithMem failed");
        return FAILED;
    }
    INFO_LOG("LoadModelFromFileWithMem success");
    ret = model_.CreateDesc();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateDesc failed");
        return FAILED;
    }
    INFO_LOG("CreateDesc success");
    ret = model_.CreateOutput();
    if (ret != SUCCESS) {
        ERROR_LOG("execute CreateOutput failed");
        return FAILED;
    }
    INFO_LOG("CreateOutput success");
    aclrtMalloc(&imageDataBuf_, (size_t)(imageDataSize_), ACL_MEM_MALLOC_HUGE_FIRST);
    if (imageDataBuf_ == nullptr) {
        ERROR_LOG("Acl malloc image buffer failed.");
        return FAILED;
    }
    INFO_LOG("aclrtMalloc success");
    ret = model_.CreateInput(imageDataBuf_, imageDataSize_);
    if (ret != SUCCESS) {
        ERROR_LOG("Create mode input dataset failed");
        return FAILED;
    }
    return SUCCESS;
}

Result RetinaObjectDetect::CreateModelInputdDataset()
{
    //Request image data memory for input model
    aclError aclRet = aclrtMalloc(&imageDataBuf_, imageDataSize_, ACL_MEM_MALLOC_HUGE_FIRST);
    if (aclRet != ACL_ERROR_NONE) {
        ERROR_LOG("malloc device data buffer failed, aclRet is %d", aclRet);
        return FAILED;
    }
    //The second input to Yolov3 is the input image width and height parameter
    const float imageInfo[4] = {(float)modelWidth_, (float)modelHeight_,
    (float)modelWidth_, (float)modelHeight_};
    imageInfoSize_ = sizeof(imageInfo);
    if (runMode_ == ACL_HOST)
        imageInfoBuf_ = Utils::CopyDataHostToDevice((void *)imageInfo, imageInfoSize_);
    else
        imageInfoBuf_ = Utils::CopyDataDeviceToDevice((void *)imageInfo, imageInfoSize_);
    if (imageInfoBuf_ == nullptr) {
        ERROR_LOG("Copy image info to device failed");
        return FAILED;
    }

    return SUCCESS;
}

Result RetinaObjectDetect::Init(int32_t& deviceId, aclrtContext& context, aclrtStream& stream, aclrtRunMode& runMode) {
    //If it is already initialized, it is returned
    if (isInited_) {
        INFO_LOG("Retina instance is initied already!");
        return SUCCESS;
    }
    Result ret;

    deviceId_ = deviceId;
    context_ = context;
    stream_ = stream;

    runMode_ = runMode;


    //Initializes the model management instance
    ret = InitModel(modelPath_);
    if (ret != SUCCESS) {
        ERROR_LOG("Init model failed");
        return FAILED;
    }
    INFO_LOG("Init model success");

    ret = (Result)dvpp_.InitResource(stream_);
    if (ret != SUCCESS) {
        ERROR_LOG("Init dvpp failed");
        return FAILED;
    }

    isInited_ = true;
    return SUCCESS;
}

// ImageData为YUV420spNV12格式的图像，Mat为bgr格式的空图像，vector为检测结果
Result RetinaObjectDetect::Process(ImageData& image, cv::Mat& bgrimg, vector<BBox>& detections) {
    // //预处理图片:读取图片,讲图片缩放到模型输入要求的尺寸
    
    struct timeval begin;
    struct timeval end;
    gettimeofday(&begin, NULL);
    Result ret = Preprocess(image, bgrimg);
    gettimeofday(&end, NULL);
    if (ret != SUCCESS) {
        ERROR_LOG("Preprocess failed");      
        return FAILED;
    }
    float time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
    INFO_LOG("Preprocess time: %fms", time_cost);

    //将预处理的图片送入模型推理,并获取推理结果
    aclmdlDataset* inferenceOutput = nullptr;
    gettimeofday(&begin, NULL);
    ret = Inference(inferenceOutput);
    gettimeofday(&end, NULL);
    if ((ret != SUCCESS) || (inferenceOutput == nullptr)) {
        ERROR_LOG("Inference model inference output data failed");
        return FAILED;
    }
    time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
    INFO_LOG("Inference time: %fms", time_cost);

    // 解析推理输出,并将推理得到的物体类别和位置标记到图片上
    gettimeofday(&begin, NULL);
    ret = PostprocessRetina(bgrimg, inferenceOutput, detections);
    gettimeofday(&end, NULL);
    if (ret != SUCCESS) {
        ERROR_LOG("Process model inference output data failed");
        return FAILED;
    }
    time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
    INFO_LOG("Postprocess time: %fms", time_cost);

    return SUCCESS;
}

Result RetinaObjectDetect::Preprocess(ImageData& srcImage, cv::Mat& bgrimg) {
//    cv::Mat inputMat;
//    cv::resize(frame, inputMat, cv::Size(modelWidth_, modelHeight_));

    // 将dvpp解码的yuv转为rgb便于后续处理
    cv::Mat yuvimg(srcImage.height * 3 / 2, srcImage.width, CV_8UC1, srcImage.data.get());
    cv::cvtColor(yuvimg, bgrimg, cv::COLOR_YUV2BGR_NV12);

    // 用dvpp做缩放，然后再用opencv做letterbox，图片越大，用这种方式的速度优势就越明显，但是dvpp是会带了一定的精度下降的
    ImageData resizedImageData;
    dvpp_.ResizeWithLetterBox(resizedImageData, srcImage, cv::Size(modelWidth_, modelHeight_));  // 用dvpp缩放
    cv::Mat resizedImage(resizedImageData.height * 3 / 2, resizedImageData.width, CV_8UC1, resizedImageData.data.get());
    cv::cvtColor(resizedImage, resizedImage, cv::COLOR_YUV2BGR_NV12);
    // cv::Mat letterBoxImage = Utils::letterbox_img(resizedImage, cv::Size(modelWidth_, modelHeight_)); // 将图像等比例缩放为模型输入大小
    aclrtMemcpyKind policy = (runMode_ == ACL_HOST)?
                             ACL_MEMCPY_HOST_TO_DEVICE:ACL_MEMCPY_DEVICE_TO_DEVICE;
    aclError ret = aclrtMemcpy(imageDataBuf_, imageDataSize_,
    resizedImage.ptr<uint8_t>(), imageDataSize_, policy);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("Copy resized image data to device failed.");
        return FAILED;
    }
    // INFO_LOG("Preprocess success");
    return SUCCESS;
}

Result RetinaObjectDetect::Inference(aclmdlDataset*& inferenceOutput) {
    //Perform reasoning
    Result ret = model_.Execute();
    if (ret != SUCCESS) {
        ERROR_LOG("Execute model inference failed");
        return FAILED;
    }
    //Get inference output
    // INFO_LOG("Inference success");
    inferenceOutput = model_.GetModelOutputData();

    return SUCCESS;
}

bool SortIm(float a, float b)
{
    return a<b;
}

Rect RetinaObjectDetect::Regression(cv::Mat &frame, float* deltas, int i)
{
    float mean[4] = {0,0,0,0};
    float std[4] = {0.1,0.1,0.2,0.2};

    float widthScale = (float)(frame.cols) / modelWidth_;
    float heightScale = (float)(frame.rows) / modelHeight_;
    Rect pred_box;
    float width = anchors[i][2] - anchors[i][0];
    float height = anchors[i][3] - anchors[i][1];
    float ctr_x = anchors[i][0] + 0.5 * width;
    float ctr_y = anchors[i][1] + 0.5 * height;

    float dx = deltas[0] * std[0] + mean[0];
    float dy = deltas[1] * std[1] + mean[1];
    float dw = deltas[2] * std[2] + mean[2];
    float dh = deltas[3] * std[3] + mean[3];

    float pred_ctr_x = ctr_x + dx * width;
    float pred_ctr_y = ctr_y + dy * height;
    float pred_w = expf(dw) * width;
    float pred_h = expf(dh) * height;

    pred_box.ltX = pred_ctr_x - 0.5 * pred_w;
    pred_box.ltY = pred_ctr_y - 0.5 * pred_h;
    pred_box.rbX = pred_ctr_x + 0.5 * pred_w;
    pred_box.rbY = pred_ctr_y + 0.5 * pred_h;

    pred_box.ltX = pred_box.ltX * widthScale;
    pred_box.ltY = pred_box.ltY * heightScale;
    pred_box.rbX = pred_box.rbX * widthScale;
    pred_box.rbY = pred_box.rbY * heightScale;
//
    pred_box.ltX = min((float)frame.cols, max((float)0.0, pred_box.ltX));
    pred_box.ltY = min((float)frame.rows, max((float)0.0, pred_box.ltY));
    pred_box.rbX = min((float)frame.cols, max((float)0.0, pred_box.rbX));
    pred_box.rbY = min((float)frame.rows, max((float)0.0, pred_box.rbY));
    
    return pred_box;
}


Result RetinaObjectDetect::PostprocessRetina(cv::Mat& frame, aclmdlDataset* modelOutput, vector<BBox>& results)
{
    uint32_t dataSize = 0;
    float* detectScore = (float*)GetInferenceOutputItem(dataSize, modelOutput,
                                            kScoreBufId);
    if (detectScore == nullptr) return FAILED;

    float* detectConf = (float*)GetInferenceOutputItem(dataSize, modelOutput,
                                              kClassBufId);
    if (detectConf == nullptr) return FAILED;

    float* detectBBox = (float*)GetInferenceOutputItem(dataSize, modelOutput,
                                              kBBoxBufId);
    if (detectBBox == nullptr) return FAILED;

   std::vector<BBox> bBox;
   double scoreThreshold = kScoreThreshold;
   for(size_t i = 0; i < kModelOutputBoxNum; i++)
   {
       float maxScore = 0;
        int maxIndex = 0;
        BBox b;

        for(int j = 0; j < 3; j++)
        {
            if(detectScore[i*3 + j] > maxScore)
            {
                maxScore = detectScore[i*3 + j];
                maxIndex = j;
                maxScore *= detectConf[i];
            }
        }

       if(maxScore > scoreThreshold)
        {
            b.score = maxScore;// * detectScore[i];
            b.cls_id = maxIndex;
        //    b.text = labels_list[maxScore];
            auto rect = Regression(frame, detectBBox + 4*i, i);
            b.rect.x = rect.ltX;
            b.rect.y = rect.ltY;
            b.rect.width = rect.rbX - rect.ltX;
            b.rect.height = rect.rbY - rect.ltY;
            bBox.push_back(b);
        }
   }

   IoU_nms(bBox, results);
//    if(!results.empty()){
//        DrawBoundBoxToImage(results, frame);
//    }
    return SUCCESS;
}

bool SortScore(BBox box1, BBox box2) {
    return box1.score < box2.score;
}

void RetinaObjectDetect::IoU_nms(std::vector<BBox> &boxes, std::vector<BBox> &result)
{
    result.clear();
    std::sort(boxes.begin(), boxes.end(), SortScore);

    while (boxes.size() != 0) {
        result.push_back(boxes.back());
        boxes.pop_back();
        int k = 0;
        size_t len = boxes.size();
        for (size_t i = 0; i < len; i++){
            float iou_value = Cal_IoU(result.back(),boxes[k]);
//            cout<<iou_value<<endl;
            if (iou_value > kNMSThreshold){
                boxes.erase((boxes.begin()+k));
            }
            else{
                k++;
            }
        }
    }
}

extern string imageName;

void RetinaObjectDetect::DrawBoundBoxToImage(const vector <BBox> &result, cv::Mat imageInput)
{
    // mtxImageId.lock();
    // uint32_t image_id = imageID.front();
    // mtxImageId.unlock();
    // string ImageFile = imageName.substr(imageName.find_last_of("/") + 1);
    // string resultFile = ImageFile.substr(0, ImageFile.find_last_of("."));
    // resultFile = "../results/" + resultFile + ".txt";
    // fstream f(resultFile, ios::out);

    for(size_t i = 0; i < result.size(); i++)
    {
        cv::Point p1, p2;
        p1.x = result[i].rect.x ;
        p1.y = result[i].rect.y ;
        p2.x = result[i].rect.x + result[i].rect.width ;
        p2.y = result[i].rect.y + result[i].rect.height ;

        cv::rectangle(imageInput, p1, p2, kColors[i % kColors.size()], kLineSolid);
        string className = labels_list[result[i].cls_id] + to_string((int)(result[i].score*100)) + "%";
        cv::putText(imageInput, className, cv::Point(p1.x, p1.y + kLabelOffset),
        cv::FONT_HERSHEY_COMPLEX, kFountScale, kFontColor);
        // f<<p1.x<<" "<<p1.y<<" "<<p2.x<<" "<<p2.y<<endl;
        printf("%d %d %d %d %s\n", p1.x, p1.y, p2.x, p2.y,className.c_str());
    }
    // f.close();

    // static int count = 0;
    // string outname = ImageFile;
    // cv::imwrite(outname, imageInput);
    // count++;
}

float RetinaObjectDetect::Cal_IoU(const BBox &b1, const BBox &b2)
{
    float w1 = b1.rect.width + 1;
    float h1 = b1.rect.height + 1;
    float w2 = b2.rect.width + 1;
    float h2 = b2.rect.width + 1;
    float area1 = w1 * h1;
    float area2 = w2 * h2;
    int x11 = std::max(b1.rect.x, b2.rect.x);
    int y11 = std::max(b1.rect.y, b2.rect.y);
    int x22 = std::min(b1.rect.x + b1.rect.width, b2.rect.x + b2.rect.width);
    int y22 = std::min(b1.rect.y + b1.rect.height, b2.rect.y + b2.rect.height);
    float intersection = (x22 - x11 + 1) * (y22 - y11 + 1);
    return intersection / (area1 + area2 - intersection);

}

void* RetinaObjectDetect::GetInferenceOutputItem(uint32_t& itemDataSize,
                                           aclmdlDataset* inferenceOutput,
                                           uint32_t idx) {

    aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(inferenceOutput, idx);
    if (dataBuffer == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer from model "
        "inference output failed", idx);
        return nullptr;
    }

    void* dataBufferDev = aclGetDataBufferAddr(dataBuffer);
    if (dataBufferDev == nullptr) {
        ERROR_LOG("Get the %dth dataset buffer address "
        "from model inference output failed", idx);
        return nullptr;
    }

    size_t bufferSize = aclGetDataBufferSizeV2(dataBuffer);
 //   cout<<(int)bufferSize<<endl;
    if (bufferSize == 0) {
        ERROR_LOG("The %dth dataset buffer size of "
        "model inference output is 0", idx);
        return nullptr;
    }

    void* data = nullptr;
    if (runMode_ == ACL_HOST) {
        data = Utils::CopyDataDeviceToLocal(dataBufferDev, bufferSize);
        if (data == nullptr) {
            ERROR_LOG("Copy inference output to host failed");
            return nullptr;
        }
    }
    else {
        data = dataBufferDev;
    }

    itemDataSize = bufferSize;
    return data;
}

void RetinaObjectDetect::DestroyResource()
{
    dvpp_.DestroyResource();
	aclrtFree(imageDataBuf_);
    aclrtFree(imageInfoBuf_);

    //The ACL resource held by the model instance must be released before the ACL exits or ABORT will be torn down
    model_.DestroyResource();

}
